#include <iostream>
#include <cmath>
#include "mpi.h"
#include "alignment.h"
#include "MbMatrix.h"
#include "model.h"
#include "settings.h"
#include "tree.h"
#include "util.h"

using namespace std;



Model::Model(Alignment *ap, MbRandom *rp, Settings *sp) {

        // Timer stuff
	double start,finish;
	double time;

	alignmentPtr = ap;
	ranPtr       = rp;
	settingsPtr  = sp;
	
	/* set up the conditional likelihoods */
	int numTaxa      = alignmentPtr->getNumTaxa();
	    numNodes     = 2 * numTaxa - 2;
	int numChar      = alignmentPtr->getNumChar();
	int numStates    = 4;
	int condLikeSize = numNodes * numChar * numStates;

	int numprocs	 = MPI::COMM_WORLD.Get_size();
	int myid		 = MPI::COMM_WORLD.Get_rank();
	int *node_off		 = new int[numprocs+1];
	int *n_node_split	 = new int[numprocs];
	split_nodes(condLikeSize, numprocs, node_off, n_node_split);
	int local_cLSize = n_node_split[myid];
	double *local_cls = new double[local_cLSize];
	cls = new double[condLikeSize];

	if ( !cls )
		{
		cerr << "ERROR: Problem allocating conditional likelihoods" << endl;
		exit(1);
		}
	
	//Using MPI Scatter and Gather to assign array
	MPI::COMM_WORLD.Scatterv(cls, n_node_split, node_off, MPI::DOUBLE, local_cls, 
				 local_cLSize, MPI::DOUBLE, 0);
	start = MPI::Wtime();
	for (int i=0; i<local_cLSize; i++)
		local_cls[i] = 0.0;
	finish = MPI::Wtime();
	time = (finish - start);
	cout << "Initialze cls time " <<time << endl;
	MPI::COMM_WORLD.Gatherv(local_cls, local_cLSize, MPI::DOUBLE, cls, n_node_split, 
				node_off, MPI::DOUBLE, 0);
	
	/* allocate a matrix of pointers so we can quickly position ourselves in
	   the vector of conditional likelihoods (cls) */
	start  = MPI::Wtime();
	clsPtr = new double**[numNodes];
	clsPtr[0] = new double*[numNodes * numChar];
	for (int i=1; i<numNodes; i++)
		clsPtr[i] = clsPtr[i-1] + numChar;
	for (int i=0; i<numNodes; i++)
		for (int j=0; j<numChar; j++)
			clsPtr[i][j] = &cls[i * numChar * numStates + j * numStates];
	finish = MPI::Wtime();
	time = (finish - start);
	cout << "Initialze clsPtr time " << time << endl;
		
	/* initialize the conditional likelihoods at the tips of the tree, which
	   are indexed 0, 1, ..., numTaxa-1 */
	start = MPI::Wtime();
	for (int i=0; i<numTaxa; i++)
		{
		for (int j=0; j<numChar; j++)
			{
			double *clp = clsPtr[i][j];
			int nucCode = alignmentPtr->getNucleotide(i, j);
			int nucs[4];
			alignmentPtr->getPossibleNucs(nucCode, nucs);
			for (int k=0; k<numStates; k++)
				{
				if ( nucs[k] == 1 )
					clp[k] = 1.0;
				}
			}
		}
	finish = MPI::Wtime();
	time = (finish - start);
	cout << "Initialze leaf time " << time << endl;
	
	/* make a tree */
	start  = MPI::Wtime();
	string treeStr = getLineFromFile( settingsPtr->getTreeFileName(), 0 ); // read the tree from the file
	tree = new Tree(alignmentPtr, ranPtr, treeStr);
	//tree->print();
	finish = MPI::Wtime();
	time = (finish - start);
	cout << "Initialze tree time " << time << endl;
	
	/* set up transition probability matrices according to the Jukes-Cantor (1969) model of
	   DNA substitution */
	start = MPI::Wtime();;
	P = new MbMatrix<double>[numNodes];
	for (int i=0; i<numNodes; i++)
		P[i] = MbMatrix<double>(4, 4);
	for (int n=0; n<numNodes; n++)
		{
		Node *p = tree->getDownPassNode( n );
		if (p->getAnc() != NULL)
			{
			int idx = p->getIndex();
			double v = p->getV();
			double x = exp(-(4.0/3.0)*v);
			for (int i=0; i<4; i++)
				{
				for (int j=0; j<4; j++)
					{
					if (i == j)
						P[idx][i][j] = 0.25 + 0.75 * x;
					else
						P[idx][i][j] = 0.25 - 0.25 * x;
					}
				}
			}
		}
	finish = MPI::Wtime();
	time = (finish - start);
	cout << "Setup Time for trans prob matrix is " <<time << endl;
	
}

Model::~Model(void) {

	delete [] cls;
	delete [] clsPtr[0];
	delete [] clsPtr;
	delete tree;
	
}

double Model::calculateLikelihood(void) {

  /* Calculate conditional likelihoods down the tree. Note that
     we have already calculated the transition probabilities in
     the constructor. */
  double *lnScaler = new double[alignmentPtr->getNumChar()];
  for (int c=0; c<alignmentPtr->getNumChar(); c++)
    lnScaler[c] = 0.0;
  for (int n=0; n<numNodes; n++)
    {
      Node *p = tree->getDownPassNode( n );
      if (p->getLft() != NULL && p->getRht() != NULL && p->getAnc() != NULL)
	{
	  /* deal with a two-way split */
	  if (p->getAnc()->getAnc() != NULL)
	    {
	      int lftIdx = p->getLft()->getIndex();
	      int rhtIdx = p->getRht()->getIndex();
	      int idx    = p->getIndex();
	      double *clL = clsPtr[lftIdx][0];
	      double *clR = clsPtr[rhtIdx][0];
	      double *clP = clsPtr[idx   ][0];
	      for (int c=0; c<alignmentPtr->getNumChar(); c++)
		{
		  double maxCl = 0.0;
		  for (int i=0; i<4; i++)
		    {
		      double sumL = 0.0, sumR = 0.0;
		      for (int j=0; j<4; j++)
			{
			  sumL += clL[j] * P[lftIdx][i][j];
			  sumR += clR[j] * P[rhtIdx][i][j];
			}
		      clP[i] = sumL * sumR;
		      if (clP[i] > maxCl)
			maxCl = clP[i];
		    }
		  lnScaler[c] += log(maxCl);
		  for (int i=0; i<4; i++)
		    clP[i] /= maxCl;
		  clL += 4;
		  clR += 4;
		  clP += 4;
		}
	    }
	  else
	    {
	      /* deal with the node to the left of the root node */
	      int lftIdx = p->getLft()->getIndex();
	      int rhtIdx = p->getRht()->getIndex();
	      int ancIdx = p->getAnc()->getIndex();
	      int idx    = p->getIndex();
	      double *clL = clsPtr[lftIdx][0];
	      double *clR = clsPtr[rhtIdx][0];
	      double *clA = clsPtr[ancIdx][0];
	      double *clP = clsPtr[idx   ][0];
	      for (int c=0; c<alignmentPtr->getNumChar(); c++)
		{
		  double maxCl = 0.0;
		  for (int i=0; i<4; i++)
		    {
		      double sumL = 0.0, sumR = 0.0, sumA = 0.0;
		      for (int j=0; j<4; j++)
			{
			  sumL += clL[j] * P[lftIdx][i][j];
			  sumR += clR[j] * P[rhtIdx][i][j];
			  sumA += clA[j] * P[idx   ][i][j];
			}
		      clP[i] = sumL * sumR * sumA;
		      if (clP[i] > maxCl)
			maxCl = clP[i];
		    }
		  lnScaler[c] += log(maxCl);
		  for (int i=0; i<4; i++)
		    clP[i] /= maxCl;
		  clL += 4;
		  clR += 4;
		  clA += 4;
		  clP += 4;
		}
	    }
	}
    }
  
  /* calculate likelihood */
  Node *p = tree->getRoot()->getLft();
  double *clP = clsPtr[p->getIndex()][0];
  double lnL = 0.0;
  for (int c=0; c<alignmentPtr->getNumChar(); c++)
    {
      double like = 0.0;
      for (int i=0; i<4; i++)
	like += clP[i] * 0.25;
      //cout << c << " -- " << log(like) << " " << alignmentPtr->getNumOfPattern(c) << endl;
      lnL += lnScaler[c] + log(like) * alignmentPtr->getNumOfPattern(c);
      clP += 4;
    }
  delete [] lnScaler;
  return lnL;
  
  
}

/* Split up the nodes evenly */
void Model::split_nodes(int numNodes, int n_proc, int *node_off, int *n_node_split) {
  
  int n_per;
  int idx, num, p;
  
  n_per = (numNodes + n_proc - 1) / n_proc;
  //if (n_per < MIN_FISH_PER_PROC) n_per = MIN_FISH_PER_PROC;
  
  idx = 0;
  num = n_per;
  for (p = 0; p < n_proc; ++p) {
    if (idx + num > numNodes)
      num = numNodes - idx;

    n_node_split[p] = num;
    node_off[p] = idx;
    idx += num;
  }
  node_off[p] = idx;

}
